rm(list = ls())
library(tidyverse)
library(estimatr)
library(grid)
library(ri2)
library(reshape2)
library(bandit)
library(stargazer)
setwd(dirname(rstudioapi::getActiveDocumentContext()$path))
set.seed(10027)
source('fn.R')

load('data.rdata')

mod  <- lm_robust(Y ~ arm - 1, data = data, weights = ipw)
true <- coef(mod)
iter <- 10000

if (file.exists('Output/TF3_simulation.rdata')){
  load('Output/TF3_simulation.rdata')
}else {
  simlist <- list()
  for (p in c(100, 500, 1000, 2000, 4000)){
    simlist[[paste0('n', p)]] <- simulate(periods = 10, n = round(p/10), probs = true, iter = iter, static = FALSE, ppmat = TRUE)
  }
}

sumlist <- list()
for (p in c(100, 500, 1000, 2000, 4000)){
  sumlist[[paste0('p', p)]] <- result_sum_NA(data = as.data.frame(simlist[[paste0('n', p)]][['d_fit']])) %>% select(-est, -bias)
}

result <- sumlist$p100 %>%
  left_join(sumlist$p500 %>% rename_at(vars(matches('best|rmse|coverage')), ~paste0(.x, '.3'))) %>%
  left_join(sumlist$p1000 %>% rename_at(vars(matches('best|rmse|coverage')), ~paste0(.x, '.4'))) %>%
  # left_join(sumlist$p2000 %>% rename_at(vars(matches('best|rmse|coverage')), ~paste0(.x, '.5'))) %>%
  left_join(sumlist$p4000 %>% rename_at(vars(matches('best|rmse|coverage')), ~paste0(.x, '.6'))) %>%
  select(-true)

stargazer::stargazer(result, summary = F, rownames = F)

# Compare to static with N = 4000
static <-  simulate(periods = 10, n = round(4000/10), probs = true, iter = iter, static = TRUE, ppmat = TRUE)
result_sum(as.data.frame(static$d_fit)) %>% select(-est,-bias)
